xtask\tasks\guest_test/
download_image.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::Xtask;
5use anyhow::Context;
6use clap::Parser;
7use clap::ValueEnum;
8use std::path::PathBuf;
9use std::process::Command;
10use vmm_test_images::KnownTestArtifacts;
11
12/// Download an image from Azure Blob Storage.
13///
14/// If no specific images are specified this command will download all available images.
15#[derive(Parser)]
16pub struct DownloadImageTask {
17    /// The folder to download the images to.
18    #[clap(short, long, default_value = "images")]
19    output_folder: PathBuf,
20    /// The test artifacts to download.
21    #[clap(long)]
22    artifacts: Vec<KnownTestArtifacts>,
23    /// Redownload images even if the file already exists.
24    #[clap(short, long)]
25    force: bool,
26}
27
28const STORAGE_ACCOUNT: &str = "hvlitetestvhds";
29const CONTAINER: &str = "vhds";
30
31impl Xtask for DownloadImageTask {
32    fn run(mut self, _ctx: crate::XtaskCtx) -> anyhow::Result<()> {
33        if self.artifacts.is_empty() {
34            self.artifacts = KnownTestArtifacts::value_variants().to_vec();
35        }
36
37        let filenames = self
38            .artifacts
39            .into_iter()
40            .map(|x| x.filename())
41            .collect::<Vec<_>>();
42
43        if !self.output_folder.exists() {
44            std::fs::create_dir(&self.output_folder)?;
45        }
46
47        let vhd_list = filenames.join(";");
48        run_azcopy_command(&[
49            "copy",
50            &format!("https://{STORAGE_ACCOUNT}.blob.core.windows.net/{CONTAINER}/*"),
51            self.output_folder.to_str().unwrap(),
52            "--include-path",
53            &vhd_list,
54            "--overwrite",
55            &self.force.to_string(),
56        ])?;
57
58        Ok(())
59    }
60}
61
62fn run_azcopy_command(args: &[&str]) -> anyhow::Result<Option<String>> {
63    let azcopy_cmd =
64        which::which("azcopy").context("Failed to find `azcopy`. Is AzCopy installed?")?;
65
66    let mut cmd = Command::new(azcopy_cmd);
67    cmd.args(args);
68
69    let mut child = cmd.spawn().context("Failed to run `azcopy` command.")?;
70    let exit = child
71        .wait()
72        .context("Failed to wait for 'azcopy' command.")?;
73    anyhow::ensure!(exit.success(), "azcopy command failed.");
74    Ok(None)
75}